# 2025电赛E题找A4 UV纸圆心，可以找到圆心和第三个圆圈，帧率 > 25fps。
# 控制云台可以基于中心点误差 err_center 进行 PID 控制
# @author Neucrack@sipeed & lxo@sieed 协助
# @license MIT
# @date 2025.7.30
from maix import camera, display, image, nn, app, time ,uart
import cv2 , struct
import numpy as np

################################ config #########################################
crop_padding = 12            # 裁切图时的外扩距离
rect_min_limit = 12          # 大黑边框四个点最小距离必须大于这个值才有效
std_from_white_rect = True   # 裁切标准图是裁切自A4纸内部白色部分
circle_num_points = 50       # 生成的第三个圆圈的点数量
std_res = [int(29.7 / 21 * 80), 80]  # 找中心点和圆圈的分辨率
hires_mode = True            # 高分辨模式，适合find_circle模式使用
high_res = 448               # 高分辨率模式宽高
model_path = "/root/models/25ec/model_3356.mud" # 检测黑框模型路径
find_circle = False          # 是否继续找圆
cam_buff_num = 1             # 摄像头缓冲
auto_awb = True              # 自动白平衡
awb_gain = [0.134, 0.0625, 0.0625, 0.1139]  # 手动白平衡值
contrast = 80                # 对比度
###################################################################################

# 初始化模型和摄像头
detector = nn.YOLOv5(model=model_path, dual_buff=True)
cam = camera.Camera(high_res, high_res, detector.input_format(), buff_num=cam_buff_num)
device1 = "/dev/ttyS0"          # 示例的串口名称                        此串口设备名称为串口0
uart1 = uart.UART(device1 , 115200)

def uart_send(serial:uart.UART,head,data):
    send_data = struct.pack(f"<B{len(data)}h",head,*data)
    serial.write(send_data)


if not auto_awb:
    cam.awb_mode(camera.AwbMode.Manual)	
    cam.set_wb_gain(awb_gain)
cam.constrast(contrast)

disp = display.Display()

err_center = [0, 0]  # 距离中心的误差
center_pos = [cam.width() // 2, cam.height() // 2]  # 画面的中心
last_center = center_pos  # 上一次检测到的圆心距离

while not app.need_exit():
    img = cam.read()
    detected_rect_area = 0
    # AI检测外框
    img_ai = img.resize(detector.input_width(), detector.input_height())
    objs = detector.detect(img_ai, conf_th=0.5, iou_th=0.45)
    
    max_idx = -1
    max_s = 0
    for i, obj in enumerate(objs):
        s = obj.w * obj.h
        if s > max_s:
            max_s = s
            max_idx = i
    
    if max_idx >= 0:
        obj = objs[max_idx]
        w = obj.w + crop_padding * 2
        h = obj.h + crop_padding * 2
        w = w + 1 if w % 2 != 0 else w
        h = h + 1 if h % 2 != 0 else h
        x = obj.x - crop_padding
        y = obj.y - crop_padding
        
        # 边界检查
        if x < 0:
            w += x
            x = 0
        if y < 0:
            h += y
            y = 0
        if x + w > img_ai.width():
            w = img_ai.width() - x
        if y + h > img_ai.height():
            h = img_ai.height() - y
        
        crop_ai = img_ai.crop(x, y, w, h)
        crop_ai_rect = [x, y, w, h]
        
        # 计算缩放比例
        img_ai_scale = [img.width() / img_ai.width(), img.height() / img_ai.height()]
        crop_rect = [int(obj.x * img_ai_scale[0]), int(obj.y * img_ai_scale[1]), 
                     int(obj.w * img_ai_scale[0]), int(h * img_ai_scale[0])]
        
        img_cv = image.image2cv(img, False, False)
        crop_ai_cv = image.image2cv(crop_ai, False, False)
        
        # 转换为灰度图
        gray = crop_ai.to_format(image.Format.FMT_GRAYSCALE)
        gray_cv = image.image2cv(gray, False, False)
        
        # 自适应二值化
        binary = cv2.adaptiveThreshold(gray_cv, 255, cv2.ADAPTIVE_THRESH_MEAN_C,
                                       cv2.THRESH_BINARY_INV, 27, 31)
        
        # 执行洪泛填充找出内白色轮廓
        if std_from_white_rect:
            h, w = binary.shape[:2]
            mask = np.zeros((h + 2, w + 2), np.uint8)
            seed_point = (2, 2)
            seed_point2 = (w - 2, h - 2)
            fill_value = 255
            cv2.floodFill(binary, mask, seed_point, fill_value, loDiff=5, upDiff=5, flags=4)
            cv2.floodFill(binary, mask, seed_point2, fill_value, loDiff=5, upDiff=5, flags=4)
            binary = cv2.bitwise_not(binary)
        
        # 查找轮廓4个角点
        contours, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        if len(contours) > 0:
            largest_contour = max(contours, key=cv2.contourArea)
            #detected_rect_area = cv2.contourArea(largest_contour)
            epsilon = 0.02 * cv2.arcLength(largest_contour, True)
            approx = cv2.approxPolyDP(largest_contour, epsilon, True)
            
            if len(approx) == 4:
                corners = approx.reshape((4, 2))
                rect = np.zeros((4, 2), dtype="float32")
                s = corners.sum(axis=1)
                rect[0] = corners[np.argmin(s)]  # 左上
                rect[2] = corners[np.argmax(s)]  # 右下
                diff = np.diff(corners, axis=1)
                rect[3] = corners[np.argmax(diff)]  # 左下
                rect[1] = corners[np.argmin(diff)]  # 右上
                
                minW = min(rect[1][0] - rect[0][0], rect[2][0] - rect[3][0])
                minH = min(rect[3][1] - rect[0][1], rect[2][1] - rect[1][1])
                
                if minH > rect_min_limit and minW > rect_min_limit:
                    maxWidth = std_res[0]
                    maxHeight = std_res[1]
                    
                    # 映射到大图
                    rect[:, 0] += crop_ai_rect[0]
                    rect[:, 1] += crop_ai_rect[1]
                    rect[:, 0] *= img_ai_scale[0]
                    rect[:, 1] *= img_ai_scale[1]
                    
                    # 透视变换
                    dst = np.array([[0, 0], [maxWidth - 1, 0], 
                                   [maxWidth - 1, maxHeight - 1], [0, maxHeight - 1]], dtype="float32")
                    M = cv2.getPerspectiveTransform(rect, dst)
                    M_inv = np.linalg.inv(M)
                    img_std_cv = cv2.warpPerspective(img_cv, M, (maxWidth, maxHeight))
                    img_std = image.cv2image(img_std_cv, False, False)

                    
                    # 计算圆圈间距
                    if std_from_white_rect:
                        circle_dist = int(img_std.height() * 0.1149425287356322)
                    else:
                        circle_dist = img_std.height() * 0.09523809523809523
                    
                    if circle_dist > 0:
                        center = [img_std.width() // 2, img_std.height() // 2]
                        center_new = None
                        
                        # 霍夫变换找圆
                        if find_circle:
                            img_std_gray_cv = cv2.cvtColor(img_std_cv, cv2.COLOR_RGB2GRAY)
                            w = h = int(circle_dist * 3)
                            roi = [center[0] - w // 2, center[1] - h // 2, w, h]
                            img_small_circle_cv = img_std_gray_cv[roi[1]:roi[1] + roi[3], roi[0]:roi[0] + roi[2]]
                            
                            circles = cv2.HoughCircles(img_small_circle_cv, cv2.HOUGH_GRADIENT, dp=1.2,
                                                    minDist=roi[2] // 2, param1=100, param2=20,
                                                    minRadius=roi[2] // 4, maxRadius=roi[2] // 2)
                            
                            if circles is not None:
                                circles = np.uint16(np.around(circles))
                                for c in circles[0, :]:
                                    center_new = (c[0] + roi[0], c[1] + roi[1])
                                    circle_dist = c[2]
                                    break
                        
                        # 更新中心点位置
                        if (not find_circle) or center_new:
                            if center_new:
                                center = center_new
                            
                            std_center_points = np.array([[center]], dtype=np.float32)
                            original_center_point = cv2.perspectiveTransform(std_center_points, M_inv)[0][0].astype(np.int32).tolist()
                            err_center = [original_center_point[0] - center_pos[0], 
                                         original_center_point[1] - center_pos[1]]
                            last_center = original_center_point
                            
                            # 计算第三个圆圈轮廓点
                            radius = circle_dist * 3
                            angles = np.linspace(0, 2 * np.pi, circle_num_points, endpoint=False)
                            cos_vals = np.cos(angles)
                            sin_vals = np.sin(angles)
                            x = center[0] + radius * cos_vals
                            y = center[1] + radius * sin_vals
                            circle_pts = np.stack((x, y), axis=1).astype(np.float32)
                            circle_pts = circle_pts[np.newaxis, :, :]
                            
                            # 反变换回原图
                            orig_circle_pts = cv2.perspectiveTransform(circle_pts, M_inv)
                            
                            # 绘制圆心和圆圈轮廓
                            img.draw_circle(original_center_point[0], original_center_point[1], 4, image.COLOR_RED, thickness=-1)
                            pts = np.round(orig_circle_pts[0]).astype(np.int32)
                            cv2.polylines(img_cv, [pts], isClosed=True, color=(0, 0, 255), thickness=1)
    
    # 绘制误差线
    img.draw_line(center_pos[0], center_pos[1], last_center[0], last_center[1], image.COLOR_RED, thickness=3)
    disp.show(img)
    ex = int(err_center[0]-67) # 镜头减去激光刚性连接的部分
    img.draw_string(0,0,f"fps is {int(time.fps())}")
    img.draw_string(0,40,f"ex is {ex}")
    print(f"err_center: ({err_center[0]:.1f}, {err_center[1]:.1f})")
    if detected_rect_area != 0:
        print(detected_rect_area)
    uart_send(uart1,0xaa,[ex,-int(err_center[1])])